import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader

# 大约共65万个神经元。
# C1	C1层的FeatureMap的神经元个数	55x55x48x2=290400
# C2	C2层的FeatureMap的神经元个数	27x27x128x2=186624
# C3	C3层的FeatureMap的神经元个数	13x13x192x2=64896
# C4	C4层的FeatureMap的神经元个数	13x13x192x2=64896
# C5	C5层的FeatureMap的神经元个数	13x13x128x2=43264
# FC6	FC6全连接层神经元个数	4096
# FC7	FC7全连接层神经元个数	4096
# Output layer	输出层神经元个数	1000
# 定义超参数
num_epochs = 5
batch_size = 100
learning_rate = 0.001

# 数据预处理
transform = transforms.Compose([
    transforms.Resize((224, 224)),  # 将MNIST图像从28x28调整为224x224，以适应AlexNet输入
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

# 加载MNIST数据集
train_dataset = torchvision.datasets.MNIST(root='./data', train=True,
                                           transform=transform, download=True)
test_dataset = torchvision.datasets.MNIST(root='./data', train=False,
                                          transform=transform)

# 创建数据加载器
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

# 定义AlexNet模型
class AlexNet(nn.Module):
    def __init__(self, num_classes=10):
        """
        初始化AlexNet模型。

        参数:
        num_classes (int): 分类的类别数，默认为10。
        """
        # 调用父类nn.Module的构造函数
        super(AlexNet, self).__init__()
        # 定义特征提取模块，包含卷积层、激活函数和池化层
        self.features = nn.Sequential(
            # 第一个卷积层，输入通道数为1，输出通道数为 64，卷积核大小为 11，步长为 4，填充为 2
            # 输入224 x 224 * 3的图像，其中 3是RGB的 3个通道；输出 55 x 55 * 64 的特征图；
            nn.Conv2d(1, 64, kernel_size=11, stride=4, padding=2),
            # 应用ReLU激活函数，inplace=True表示直接在原张量上操作，节省内存
            nn.ReLU(inplace=True),
            # 第一个最大池化层，池化核大小为 3，步长为 2
            # 输出27x27的特征图；
            nn.MaxPool2d(kernel_size = 3, stride = 2),
            # 第二个卷积层，输入通道数为64，输出通道数为192，卷积核大小为5，填充为2
            nn.Conv2d(64, 192, kernel_size=5, padding=2),
            # 应用ReLU激活函数
            nn.ReLU(inplace=True),
            # 第二个最大池化层，池化核大小为3，步长为2
            nn.MaxPool2d(kernel_size=3, stride=2),
            # 第三个卷积层，输入通道数为192，输出通道数为384，卷积核大小为3，填充为1
            nn.Conv2d(192, 384, kernel_size=3, padding=1),
            # 应用ReLU激活函数
            nn.ReLU(inplace=True),
            # 第四个卷积层，输入通道数为384，输出通道数为256，卷积核大小为3，填充为1
            nn.Conv2d(384, 256, kernel_size=3, padding=1),
            # 应用ReLU激活函数
            nn.ReLU(inplace=True),
            # 第五个卷积层，输入通道数为256，输出通道数为256，卷积核大小为3，填充为1
            nn.Conv2d(256, 256, kernel_size=3, padding=1),
            # 应用ReLU激活函数
            nn.ReLU(inplace=True),
            # 第三个最大池化层，池化核大小为3，步长为2
            nn.MaxPool2d(kernel_size=3, stride=2),
        )
        # 定义自适应平均池化层，将输出特征图的大小调整为6x6
        self.avgpool = nn.AdaptiveAvgPool2d((6, 6))
        # 定义分类器模块，包含全连接层、激活函数和Dropout层
        self.classifier = nn.Sequential(
            # 第一个Dropout层，用于防止过拟合
            nn.Dropout(),
            # 第一个全连接层，输入特征数为 256 * 6 * 6 = 9216，输出特征数为4096
            nn.Linear(256 * 6 * 6, 4096),
            # 应用ReLU激活函数
            nn.ReLU(inplace=True),
            # 第二个Dropout层
            nn.Dropout(),
            # 第二个全连接层，输入特征数为4096，输出特征数为4096
            nn.Linear(4096, 4096),
            # 应用ReLU激活函数
            nn.ReLU(inplace=True),
            # 输出全连接层，输入特征数为4096，输出特征数为分类的类别数
            nn.Linear(4096, num_classes),
        )

    def forward(self, x):
        x = self.features(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        return x

# 实例化模型
model = AlexNet()

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)

# 训练模型
total_step = len(train_loader)
for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        # 前向传播
        outputs = model(images)
        loss = criterion(outputs, labels)

        # 反向传播和优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if (i + 1) % 1 == 0:
            print(f'Epoch [{epoch + 1}/{num_epochs}], Step [{i + 1}/{total_step}], Loss: {loss.item():.4f}')

# 测试模型
model.eval()
with torch.no_grad():
    correct = 0
    total = 0
    for images, labels in test_loader:
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

    print(f'Accuracy of the network on the 10000 test images: {100 * correct / total}%')
